{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib qt\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "from local_plot import *\n",
    "from utils import *\n",
    "#Plot the IMU data\n",
    "from unicodedata import name\n",
    "from scipy.fft import fft, fftfreq\n",
    "import matplotlib\n",
    "matplotlib.rcParams.update({'font.size': 30})\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import rosbag\n",
    "def read_imu(bagpath, start=0.0):\n",
    "    ts = []\n",
    "    acc = []\n",
    "    gyro = []\n",
    "    t0 = None\n",
    "    for topic, msg, t in rosbag.Bag(bagpath).read_messages():\n",
    "        if msg._type == \"sensor_msgs/Imu\":\n",
    "            if t0 is None:\n",
    "                t0 = msg.header.stamp.to_sec()\n",
    "            if msg.header.stamp.to_sec() - t0 < start:\n",
    "                continue\n",
    "            ts.append(msg.header.stamp.to_sec() - t0)\n",
    "            acc.append([msg.linear_acceleration.x, msg.linear_acceleration.y, msg.linear_acceleration.z])\n",
    "            gyro.append([msg.angular_velocity.x, msg.angular_velocity.y, msg.angular_velocity.z])\n",
    "    imu = {}\n",
    "    imu[\"t\"], imu[\"acc\"], imu[\"gyro\"] = np.array(ts), np.array(acc), np.array(gyro)\n",
    "    return imu\n",
    "\n",
    "# imu = read_imu(\"/home/xuhao/data/d2slam/manual_vicon_quad_2022_10_25_small/manual_small_1-sync.bag\", 20)\n",
    "# bagname = \"/home/xuhao/data/d2slam/manual_vicon_quad_2022_10_25_small/manual_data_calib-sync.bag\"\n",
    "bagname = \"/home/xuhao/data/d2slam/manual_vicon_quad_2022_10_25_small/manual_small_6-sync.bag\"\n",
    "imu = read_imu(bagname)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fft_data(data, fs):\n",
    "    N = len(data)\n",
    "    T = 1.0 / fs\n",
    "    yf = fft(data)\n",
    "    xf = fftfreq(N, T)[:N//2]\n",
    "    return xf, 2.0/N * np.abs(yf[0:N//2])\n",
    "def draw_fft_data(data, fs, label, ax, axis=\"\"):\n",
    "    xf, yf = fft_data(data, fs)\n",
    "    ax.semilogy(xf, yf, label=label)\n",
    "    ax.set_ylabel(f\"Amp. {axis}\")\n",
    "    # ax.set_ylim(0, 1.2*np.max(yf))\n",
    "    ax.legend()\n",
    "    ax.grid()\n",
    "def draw_fft(imu, name, item=\"acc\", axs=None, figsize=(20,10)):\n",
    "    if axs is None:\n",
    "        plt.figure(f\"{name}_fft_{item}\", figsize=figsize)\n",
    "        plt.clf()\n",
    "        fig, axs = plt.subplots(3,1, sharex=True, figsize=figsize, num=f\"{name}_fft_{item}\")\n",
    "    draw_fft_data(imu[item][:,0], 220.0, name, axs[0], \"x\")\n",
    "    draw_fft_data(imu[item][:,1], 220.0, name, axs[1], \"y\")\n",
    "    draw_fft_data(imu[item][:,2], 220.0, name, axs[2], \"z\")\n",
    "    axs[2].set_xlabel(\"Frequency (Hz)\")\n",
    "    plt.tight_layout()\n",
    "    return axs\n",
    "    \n",
    "def draw_imu_data(imu, suffix=\"\", axs=None, axs_gyro=None, figsize = (20,10)):\n",
    "    if axs is None:\n",
    "        plt.figure(f\"acc\", figsize=figsize)\n",
    "        plt.clf()\n",
    "        fig, axs = plt.subplots(3,1, sharex=True, figsize=figsize, num=f\"acc\")\n",
    "\n",
    "        plt.figure(f\"gyro\", figsize=figsize)\n",
    "        plt.clf()\n",
    "        fig, axs_gyro = plt.subplots(3,1, sharex=True, figsize=figsize, num=f\"gyro\")\n",
    "\n",
    "    axs[0].plot(imu[\"t\"], imu[f\"acc{suffix}\"][:,0] - np.mean(imu[f\"acc{suffix}\"][:,0]), label=f\"a_x {suffix}\")\n",
    "    plt.legend()\n",
    "    axs[1].plot(imu[\"t\"], imu[f\"acc{suffix}\"][:,1] - np.mean(imu[f\"acc{suffix}\"][:,1]), label=f\"a_y {suffix}\")\n",
    "    plt.legend()\n",
    "    axs[2].plot(imu[\"t\"], imu[f\"acc{suffix}\"][:,2], label=f\"a_z {suffix}\")\n",
    "    plt.legend()\n",
    "    axs_gyro[0].plot(imu[\"t\"], imu[f\"gyro{suffix}\"][:,0], label=\"w_x\")\n",
    "    plt.legend()\n",
    "    axs_gyro[1].plot(imu[\"t\"], imu[f\"gyro{suffix}\"][:,1], label=\"w_y\")\n",
    "    plt.legend()\n",
    "    axs_gyro[2].plot(imu[\"t\"], imu[f\"gyro{suffix}\"][:,2], label=\"w_z\")\n",
    "    return axs, axs_gyro\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Appy 2nd order butterworth filter\n",
    "from matplotlib.lines import lineStyles\n",
    "from scipy.signal import butter, lfilter, iirnotch, filtfilt\n",
    "def butter_lowpass_filter(data, cutoff, fs, order=5):\n",
    "    nyq = 0.5 * fs\n",
    "    normal_cutoff = cutoff / nyq\n",
    "    b, a = butter(order, normal_cutoff, btype='low', analog=False)\n",
    "    y = lfilter(b, a, data)\n",
    "    return y\n",
    "#Apply a notch filter\n",
    "def notch_filter(data, cut_freq, fs, Q = 30.0):\n",
    "    b, a = iirnotch(cut_freq, Q, fs)\n",
    "    y = filtfilt(b, a, data)\n",
    "    return y\n",
    "#Apply 2nd order butterworth filter on the data\n",
    "acc = imu[\"acc\"]\n",
    "gyro = imu[\"gyro\"]\n",
    "fs = 220.0\n",
    "cutoff = 20.0\n",
    "order = 2\n",
    "acc_filtered = np.zeros(acc.shape)\n",
    "gyro_filtered = np.zeros(gyro.shape)\n",
    "for i in range(3):\n",
    "    acc_filtered[:,i] = butter_lowpass_filter(acc[:,i], cutoff, fs, order)\n",
    "    gyro_filtered[:,i] = butter_lowpass_filter(gyro[:,i], cutoff, fs, order)\n",
    "imu[\"acc_lowpass\"] = acc_filtered\n",
    "imu[\"gyro_lowpass\"] = gyro_filtered\n",
    "#Plot the filtered data with raw\n",
    "cut_freq = 65\n",
    "Q = 3\n",
    "acc_filtered = np.zeros(acc.shape)\n",
    "gyro_filtered = np.zeros(gyro.shape)\n",
    "for i in range(3):\n",
    "    acc_filtered[:,i] = notch_filter(acc[:,i], cut_freq, fs, Q)\n",
    "    gyro_filtered[:,i] = notch_filter(gyro[:,i], cut_freq, fs, Q)\n",
    "imu[\"acc_notch\"] = acc_filtered\n",
    "imu[\"gyro_notch\"] = gyro_filtered\n",
    "\n",
    "acc_filtered = np.zeros(acc.shape)\n",
    "gyro_filtered = np.zeros(gyro.shape)\n",
    "for i in range(3):\n",
    "    acc_filtered[:,i] = butter_lowpass_filter(imu[\"acc_notch\"][:,i], cutoff, fs, order)\n",
    "    gyro_filtered[:,i] = butter_lowpass_filter(imu[\"gyro_notch\"][:,i], cutoff, fs, order)\n",
    "imu[\"acc_notch_lowpass\"] = acc_filtered\n",
    "imu[\"gyro_notch_lowpass\"] = gyro_filtered\n",
    "\n",
    "figsize = (20,15)\n",
    "axs, axs_gyro = draw_imu_data(imu, suffix=\"\", figsize=figsize)\n",
    "# axs = draw_imu_data(imu, suffix=\"_butter\", axs=axs)\n",
    "axs, axs_gyro = draw_imu_data(imu, suffix=\"_notch\", axs=axs, axs_gyro=axs_gyro)\n",
    "axs, axs_gyro = draw_imu_data(imu, suffix=\"_lowpass\", axs=axs, axs_gyro=axs_gyro)\n",
    "axs, axs_gyro = draw_imu_data(imu, suffix=\"_notch_lowpass\", axs=axs, axs_gyro=axs_gyro)\n",
    "\n",
    "axs = draw_fft(imu, \"raw\", \"acc\", figsize=figsize)\n",
    "axs = draw_fft(imu, \"notch\", item=\"acc_notch\", axs=axs)\n",
    "axs = draw_fft(imu, \"notch_lp\", item=\"acc_notch_lowpass\", axs=axs)\n",
    "\n",
    "axs_gyro = draw_fft(imu, \"raw\", \"gyro\", figsize=figsize)\n",
    "axs_gyro = draw_fft(imu, \"notch\", item=\"gyro_notch\", axs=axs_gyro)\n",
    "axs_gyro = draw_fft(imu, \"lowpass\", item=\"gyro_lowpass\", axs=axs_gyro)\n",
    "axs_gyro = draw_fft(imu, \"notch_lp\", item=\"gyro_notch_lowpass\", axs=axs_gyro)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def write_imu_filtered(baginput, imu, start=0.0, acc_suffix=\"_notch\", gyro_suffix=\"_lowpass\"):\n",
    "    ts = []\n",
    "    acc = []\n",
    "    gyro = []\n",
    "    t0 = None\n",
    "    #baginput is string, convert a output name\n",
    "    bagoutput = baginput.replace(\".bag\", f\"_filtered.bag\")\n",
    "    c = 0\n",
    "    with rosbag.Bag(bagoutput, 'w') as outbag:\n",
    "        for topic, msg, t in rosbag.Bag(baginput).read_messages():\n",
    "            if msg._type == \"sensor_msgs/Imu\":\n",
    "                if t0 is None:\n",
    "                    t0 = msg.header.stamp.to_sec()\n",
    "                if msg.header.stamp.to_sec() - t0 < start:\n",
    "                    continue\n",
    "                #Use filtered imu data to replace\n",
    "                msg.linear_acceleration.x = imu[\"acc\"+acc_suffix][c,0]\n",
    "                msg.linear_acceleration.y = imu[\"acc\"+acc_suffix][c,1]\n",
    "                msg.linear_acceleration.z = imu[\"acc\"+acc_suffix][c,2]\n",
    "                msg.angular_velocity.x = imu[\"gyro\"+gyro_suffix][c,0]\n",
    "                msg.angular_velocity.y = imu[\"gyro\"+gyro_suffix][c,1]\n",
    "                msg.angular_velocity.z = imu[\"gyro\"+gyro_suffix][c,2]\n",
    "                c += 1\n",
    "                outbag.write(topic, msg, t)\n",
    "            else:\n",
    "                outbag.write(topic, msg, t)\n",
    "write_imu_filtered(bagname, imu, start=0.0, acc_suffix=\"_notch_lowpass\", gyro_suffix=\"_notch_lowpass\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.5 ('base')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "08ce52785f0fedc81003ce387e097a83d6cc9494681cd746006386992005bb71"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
